package com.demo.config.sharding;

import com.google.common.collect.Range;
import com.google.common.collect.Sets;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingValue;

import java.util.Collection;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

/**
 * 复合分片策略
 * @author K
 */
@Slf4j
public class ComplexKeysShardingStrategy implements ComplexKeysShardingAlgorithm {

    @Override
    public Collection<String> doSharding(Collection availableTargetNames, ComplexKeysShardingValue shardingValue) {
        log.info("强制分配分片策略，表:{}，ShardingValue:{}", availableTargetNames, shardingValue);

        Map<String, Collection> columnNameAndShardingValuesMap = shardingValue.getColumnNameAndShardingValuesMap();
        Map<String, Range> columnNameAndRangeValuesMap = shardingValue.getColumnNameAndRangeValuesMap();

        for (String column : ShardingConstant.COLUMNS) {
            if (columnNameAndShardingValuesMap.containsKey(column)) {
                Collection shardingValues = columnNameAndShardingValuesMap.get(column);
                return preciseDb(availableTargetNames, shardingValues);
            }

            if (columnNameAndRangeValuesMap.containsKey(column)) {
                Range range = columnNameAndRangeValuesMap.get(column);
                return rangDb(availableTargetNames, range);
            }
        }

        return availableTargetNames;
    }

    private Set<String> preciseDb(Collection availableTargetNames, Collection shardingValues) {
        Set<String> database = new HashSet<>();
        for (Object obj : shardingValues) {
            Long mod = null;
            if (obj instanceof Long) {
                Long tmp = (Long) obj;
                mod = tmp % ShardingConstant.DB_NUM;
            } else {
                return Sets.newHashSet(availableTargetNames);
            }

            String tmpDbName = ShardingConstant.DB_NAME + mod;
            if (availableTargetNames.contains(tmpDbName)) {
                database.add(tmpDbName);
            }
        }
        log.info("分片后数据库：{}", database);
        return database;
    }

    private Set<String> rangDb(Collection availableTargetNames, Range range) {
        Long lowerEndpoint = range.hasLowerBound() ? (Long) range.lowerEndpoint() : 0L;
        Long upperEndpoint = range.hasUpperBound() ? (Long) range.upperEndpoint() : Long.MAX_VALUE;

        Set<String> database = new HashSet<>();
        // 根据范围值进行遍历，计算需要查询的db
        for (long tmp = lowerEndpoint; tmp <= upperEndpoint; tmp ++) {
            long mod = Math.abs(tmp % ShardingConstant.DB_NUM);
            String tmpDbName = ShardingConstant.DB_NAME + mod;

            // 集合里的db和availableTargetNames长度相同，证明已经把所有db都遍历完
            // 所以结束遍历
            if (database.size() == availableTargetNames.size()) {
                break;
            }

            if (database.contains(tmpDbName)) {
                continue;
            }

            if (availableTargetNames.contains(tmpDbName)) {
                database.add(tmpDbName);
            }
        }
        return database;
    }
}
